#include <iostream>
#include <cmath>
#include "alignment.h"
#include "MbMatrix.h"
#include "model.h"
#include "settings.h"
#include "tree.h"
#include "util.h"
//#include "mpi.h"

using namespace std;



Model::Model(Alignment *ap, MbRandom *rp, Settings *sp) {

	alignmentPtr = ap;
	ranPtr       = rp;
	settingsPtr  = sp;
	localCls = NULL;
	localClsPtr = NULL;
	subtrees = NULL;
	queue = NULL;
	
	/* set up the conditional likelihoods */
	int numTaxa      = alignmentPtr->getNumTaxa();
	    numNodes     = 2 * numTaxa - 2;
	int numChar      = alignmentPtr->getNumChar();
	int numStates    = 4;

	/* make a tree */
	string treeStr = getLineFromFile( settingsPtr->getTreeFileName(), 0 ); // read the tree from the file
	tree = new Tree(alignmentPtr, ranPtr, treeStr);
	//tree->print();
	
	/* set up transition probability matrices according to the Jukes-Cantor (1969) model of
	   DNA substitution */
	P = new MbMatrix<double>[numNodes];
	for (int i=0; i<numNodes; i++)
		P[i] = MbMatrix<double>(4, 4);
	for (int n=0; n<numNodes; n++)
		{
		Node *p = tree->getDownPassNode( n );
		if (p->getAnc() != NULL)
			{
			int idx = p->getIndex();
			double v = p->getV();
			double x = exp(-(4.0/3.0)*v);
			for (int i=0; i<4; i++)
				{
				for (int j=0; j<4; j++)
					{
					if (i == j)
						P[idx][i][j] = 0.25 + 0.75 * x;
					else
						P[idx][i][j] = 0.25 - 0.25 * x;
					}
				}
			}
		}
	
}

Model::~Model(void) {
	/*
	delete [] cls;
	delete [] clsPtr[0];
	delete [] clsPtr;
	*/
	delete tree;
	if (subtrees != NULL)
		delete [] subtrees;
	if (queue !=NULL)
		delete queue;
	if (localCls !=NULL)
		delete localCls;
	if (localClsPtr !=NULL)
		delete localClsPtr;
}

double Model::calculateLikelihood(void) {
	
		/* Split up tree */
		int numST = 4;	
		splitTree(numST);
		double *clL, *clR, *clP, *clA;
		
		/* Calculate conditional likelihoods down the tree. Note that
		   we have already calculated the transition probabilities in
		   the constructor. */
		double *lnScaler = new double[alignmentPtr->getNumChar()];
		for (int c=0; c<alignmentPtr->getNumChar(); c++)
			lnScaler[c] = 0.0;	
		
		SubTree *st = queue->dequeue();
		int idxST = st->getIndexST();
		
		while (st != NULL)
			{
			//cout << "number of node for subtree " << st->getNumNodesST() << endl;
				
			localCls = subtrees[idxST]->getClsST();
			localClsPtr = subtrees[idxST]->getClsPtrST(localCls);
			// First element in localCls is subtree index
			//idx = (int) localCls[0];
			for (int n=0; n<subtrees[idxST]->getNumNodesST(); n++)
				{
				Node *p = subtrees[idxST]->getDownPassNodeST( n );
				
				//cout << "MARKER-it's working here" << n << " " << p->getIndex() << endl;
				//cout << p->getIsChildRoot() << endl;
				if (p->getLft() != NULL && p->getRht() != NULL 
						&& p->getAnc() != NULL && !p->getIsChildRoot() && p != st->getSubRoot())
					{
					/* deal with a two-way split */
					if (p->getAnc()->getAnc() != NULL)
						{
				    int lftIdx = p->getLft()->getIndex();
				    int rhtIdx = p->getRht()->getIndex();
						// Find numNumber, n, for localClsPtr for left and right nodes
						for (int k=n-1; k>=0; k--) 
							{
							Node *q = subtrees[idxST]->getDownPassNodeST( k );
							if (lftIdx == q->getIndex()) 
								clL = localClsPtr[k];
							if (rhtIdx == q->getIndex())
								clR = localClsPtr[k];
							}
											
						clP = localClsPtr[n];
						for (int c=0; c<alignmentPtr->getNumChar(); c++)
							{
							double maxCl = 0.0;
							for (int i=0; i<4; i++)
								{
								double sumL = 0.0, sumR = 0.0;
								for (int j=0; j<4; j++)
									{	
									sumL += clL[j] * P[lftIdx][i][j];
									sumR += clR[j] * P[rhtIdx][i][j];
									}
									//cout << "Marker " << endl;
								clP[i] = sumL * sumR;
								if (clP[i] > maxCl)
									maxCl = clP[i];
								}
							lnScaler[c] += log(maxCl);
							for (int i=0; i<4; i++)
								clP[i] /= maxCl;
							clL += 4;
							clR += 4;
							clP += 4;
							
							}
						}
					else
						{
						/* deal with the node to the left of the root node */
						int lftIdx = p->getLft()->getIndex();
					  int rhtIdx = p->getRht()->getIndex();
						int idx    = p->getIndex();
						// Find numNumber, n, for localClsPtr for left and right nodes
						for (int k=n-1; k>=0; k--) 
							{
							Node *q = subtrees[idxST]->getDownPassNodeST( k );
							if (lftIdx == q->getIndex()) 
								clL = localClsPtr[k];
							if (rhtIdx == q->getIndex())
								clR = localClsPtr[k];
							}
						clA = localClsPtr[n+1];
						clP = localClsPtr[n];
						for (int c=0; c<alignmentPtr->getNumChar(); c++)
							{
							double maxCl = 0.0;
							for (int i=0; i<4; i++)
								{
								double sumL = 0.0, sumR = 0.0, sumA = 0.0;
								for (int j=0; j<4; j++)
									{
									sumL += clL[j] * P[lftIdx][i][j];
									sumR += clR[j] * P[rhtIdx][i][j];
									sumA += clA[j] * P[idx][i][j];
									}
								clP[i] = sumL * sumR * sumA;
								if (clP[i] > maxCl)
									maxCl = clP[i];
								}
							lnScaler[c] += log(maxCl);
							for (int i=0; i<4; i++)
								clP[i] /= maxCl;
							clL += 4;
							clR += 4;
							clA += 4;
							clP += 4;
							}
						}
					}
					
				if(p == st->getSubRoot() && p->getAnc() != NULL)
					{
						
					int lftIdx = p->getLft()->getIndex();
				  int rhtIdx = p->getRht()->getIndex();
						// Find numNumber, n, for localClsPtr for left and right nodes
					for (int k=n-1; k>=0; k--) 
						{
						Node *q = subtrees[idxST]->getDownPassNodeST( k );
						if (lftIdx == q->getIndex()) 
							clL = localClsPtr[k];
						if (rhtIdx == q->getIndex())
							clR = localClsPtr[k];
						}
												
						clP = localClsPtr[n];
						for (int c=0; c<alignmentPtr->getNumChar(); c++)
							{
							double maxCl = 0.0;
							for (int i=0; i<4; i++)
								{
								double sumL = 0.0, sumR = 0.0;
								for (int j=0; j<4; j++)
									{
									sumL += clL[j] * P[lftIdx][i][j];
									sumR += clR[j] * P[rhtIdx][i][j];
									}
								clP[i] = sumL * sumR;
								if (clP[i] > maxCl)
									maxCl = clP[i];
								}
							lnScaler[c] += log(maxCl);
							for (int i=0; i<4; i++)
								clP[i] /= maxCl;
							clL += 4;
							clR += 4;
							clP += 4;
						}
					}
				}
			subtrees[idxST]->setClsST(localCls);
			queue->completeSubTree(st);
			st = queue->dequeue();
			if (st != NULL)
				idxST = st->getIndexST();
			}
	/* calculate likelihood */
	Node *p = tree->getRoot()->getLft();
	clP = p->getClsPtr(0);
	double lnL = 0.0;
	for (int c=0; c<alignmentPtr->getNumChar(); c++)
		{
		double like = 0.0;
		for (int i=0; i<4; i++)
			like += clP[i] * 0.25;
		//cout << c << " -- " << log(like) << " " << alignmentPtr->getNumOfPattern(c) << endl;
		lnL += lnScaler[c] + log(like) * alignmentPtr->getNumOfPattern(c);
		clP += 4;
		}
	delete [] lnScaler;
	
	return lnL;

}

void Model::splitTree(int numSubTrees) {
	
	/* find Sub-roots for subtrees */
	int nPer, idx, candIdx;
	Node **subroots = new Node*[numSubTrees];
	Node *candidate;
	
	nPer = (numNodes + numSubTrees - 1) / numSubTrees;
	idx = nPer-1;
	
  for (int i = 0; i<numSubTrees; i++) {
		
		if (idx >= numNodes-1)
      subroots[i] = tree->getDownPassNode( numNodes - 1 );
		else {
			candidate = tree->getDownPassNode( idx );
			candIdx = idx;
			while (candidate->getLft() == NULL || candidate->getRht() == NULL)
				candidate = tree->getDownPassNode(++candIdx);
		
			subroots[i] = tree->getDownPassNode( candIdx );	
    
			idx += nPer;
			if (candIdx > idx+nPer){
				idx = candIdx;
				nPer = (numNodes - candIdx + numSubTrees - i + 1) / (numSubTrees - i + 1);
			}
		}
  }


	/*
	//Custom input for bglobin.in where numSubTrees=4 FOR TESTING ONLY
	subroots[0] = tree->getDownPassNode( 8 );
	subroots[1] = tree->getDownPassNode( 15 );
	subroots[2] = tree->getDownPassNode( 21 ); 
	subroots[3] = tree->getDownPassNode( 31 );
	*/
	
	subtrees = new SubTree*[numSubTrees];
	queue = new Queue(numSubTrees);
	
	for (int i = 0; i < numSubTrees; i++) {
		subtrees[i] = new SubTree(tree, subroots[i], i);
	}
	
	//cout << "number of children for Node " << subroots[0]->getIndex() 
	//			<< " = " <<  subroots[i]->getNumChildren() << endl;
	//cout << "numChildRoot "<< 0 << " = " << subtrees[0]->numChildRoot << endl;
	
	/*subtrees[0]->setAncSubTree(subtrees[3]);
	subtrees[1]->setAncSubTree(subtrees[3]);
	subtrees[2]->setAncSubTree(subtrees[3]);
	*/
	
	/* set AncSubTrees -- There is gotta be a better way to do this */
	for (int i = 0; i < numSubTrees; i++) {
		if (subtrees[i]->numChildRoot == 0)
			queue->enqueue(subtrees[i]);
		else {
			for (int j=0; j<subtrees[i]->getNumNodesST(); j++) {
				
				Node *LftNode = subtrees[i]->getDownPassNodeST(j)->getLft();
				Node *RhtNode = subtrees[i]->getDownPassNodeST(j)->getRht();
				
				if (LftNode != NULL && RhtNode != NULL) {
					
					if (LftNode->getIsChildRoot()) {
						for (int k=0; k<numSubTrees; k++) 
							if (LftNode->getIndex() == subtrees[k]->getSubRoot()->getIndex())
								subtrees[k]->setAncSubTree(subtrees[i]);
						}
						
					if (RhtNode->getIsChildRoot()) {
						for (int k=0; k<numSubTrees; k++)
							if (RhtNode->getIndex() == subtrees[k]->getSubRoot()->getIndex())
								subtrees[k]->setAncSubTree(subtrees[i]);
					}
				}
			}
		}
	}
	
	delete [] subroots;
	
}
